package main

import (
	"encoding/hex"
	"log"
	"math/big"

	"github.com/consensys/gnark-crypto/ecc"
	"github.com/consensys/gnark-crypto/ecc/bn254/fr/mimc"
	"github.com/consensys/gnark/backend/groth16"
	"github.com/consensys/gnark/frontend"
	"github.com/consensys/gnark/frontend/cs/r1cs"

	"client2/zkp"
)

func main() {
	log.Println("🔍 测试 MiMC 哈希计算一致性...")

	// 测试数据
	secretHex := "a173f2303c1c6ef28e04ba656704551e2332809fca8dd2a4094cd991f5ddbeeb"
	deviceID := "TEMP_SENSOR"
	challengeHex := "1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef"

	log.Printf("Secret: %s", secretHex)
	log.Printf("DeviceID: %s", deviceID)
	log.Printf("Challenge: %s", challengeHex)

	// === 方法1: 使用 prover 的 MiMC 实现（gnark-crypto） ===
	log.Println("\n=== 方法1: 使用 prover 的 MiMC 实现 ===")

	secretBytes, _ := hex.DecodeString(secretHex)
	deviceIDBytes := []byte(deviceID)
	challengeBytes, _ := hex.DecodeString(challengeHex)

	// 计算 commitment = MiMC(secret, deviceID)
	commitment1 := computeMiMCWithCrypto(secretBytes, deviceIDBytes)
	log.Printf("Commitment (crypto): %s", commitment1)

	// 计算 response = MiMC(secret, challenge)
	response1 := computeMiMCWithCrypto(secretBytes, challengeBytes)
	log.Printf("Response (crypto): %s", response1)

	// === 方法2: 通过电路计算（witness） ===
	log.Println("\n=== 方法2: 通过电路计算 witness ===")

	// 转换为 big.Int
	secretBig := new(big.Int).SetBytes(secretBytes)
	deviceIDBig := new(big.Int).SetBytes(deviceIDBytes)
	challengeBig := new(big.Int).SetBytes(challengeBytes)
	commitmentBig := new(big.Int).SetBytes(hexToBytes(commitment1))
	responseBig := new(big.Int).SetBytes(hexToBytes(response1))

	log.Printf("SecretBig: %s", secretBig.String())
	log.Printf("DeviceIDBig: %s", deviceIDBig.String())
	log.Printf("CommitmentBig: %s", commitmentBig.String())
	log.Printf("ChallengeBig: %s", challengeBig.String())
	log.Printf("ResponseBig: %s", responseBig.String())

	// === 方法3: 编译电路验证约束 ===
	log.Println("\n=== 方法3: 编译电路并生成证明 ===")

	circuit := &zkp.AuthCircuit{}
	ccs, err := frontend.Compile(ecc.BN254.ScalarField(), r1cs.NewBuilder, circuit)
	if err != nil {
		log.Fatal("编译电路失败:", err)
	}
	log.Printf("电路编译成功，约束数量: %d", ccs.GetNbConstraints())

	// 创建见证
	assignment := &zkp.AuthCircuit{
		Secret:     secretBig,
		DeviceID:   deviceIDBig,
		Challenge:  challengeBig,
		Commitment: commitmentBig,
		Response:   responseBig,
	}

	witness, err := frontend.NewWitness(assignment, ecc.BN254.ScalarField())
	if err != nil {
		log.Fatal("创建见证失败:", err)
	}
	log.Println("✅ 见证创建成功")

	// 尝试验证约束（使用 dummy setup）
	pk, vk, err := groth16.Setup(ccs)
	if err != nil {
		log.Fatal("Setup失败:", err)
	}
	log.Println("✅ Setup完成")

	// 生成证明
	proof, err := groth16.Prove(ccs, pk, witness)
	if err != nil {
		log.Printf("❌ 生成证明失败: %v", err)
		return
	}
	log.Println("✅ 证明生成成功！")

	// 验证证明
	publicWitness, _ := witness.Public()
	err = groth16.Verify(proof, vk, publicWitness)
	if err != nil {
		log.Printf("❌ 验证失败: %v", err)
		return
	}
	log.Println("✅ 证明验证成功！")
}

func computeMiMCWithCrypto(secretBytes, dataBytes []byte) string {
	// 将 dataBytes 转换为域元素（32 字节）
	dataBig := new(big.Int).SetBytes(dataBytes)
	dataFieldBytes := make([]byte, 32)
	dataBig.FillBytes(dataFieldBytes)

	// 创建 MiMC 哈希
	h := mimc.NewMiMC()
	h.Write(secretBytes)
	h.Write(dataFieldBytes)

	result := h.Sum(nil)
	return hex.EncodeToString(result)
}

func hexToBytes(hexStr string) []byte {
	bytes, _ := hex.DecodeString(hexStr)
	return bytes
}
